package org.example;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

class Solution {

    public int[] smallestK(int[] arr, int k) {
        quickSort(arr, 0, arr.length - 1, k);
        int[] ret = new int[k];
        for (int i = 0; i < k; i++) {
            ret[i] = arr[i];
        }
        return ret;
    }

    private void quickSort(int[] arr, int l, int r, int k) {
        if(l >= r) return;

        //随机数取基准值
        int key = arr[new Random().nextInt(r - l + 1) + l];

        //三路划分
        int left = l - 1, right = r + 1, i = l;
        while(i < right){
            if(arr[i] < key) swap(arr, ++left, i++);
            else if(arr[i] == key) i++;
            else swap(arr, --right, i);
        }
        //数组被划分为三个区间
        // [l, left] [left + 1, right - 1] [right, r]
        //计算三个区间的元素个数 a, b, c
        int a = left - l + 1, b = right - left - 1;
        if(a >= k) quickSort(arr, l, left, k);
        else if(a + b >= k) return;
        else quickSort(arr, right, r, k - a - b);
    }

    private void swap(int[] arr, int i, int j) {
        int t = arr[i];
        arr[i] = arr[j];
        arr[j] = t;
    }
}
